import os
from dotenv import load_dotenv
# 之前的引用: from langchain_openai import ChatOpenAI, OpenAIEmbeddings
# 新的引用:
from langchain_community.chat_models import ChatOllama
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

from http import HTTPStatus
from dashscope import ImageSynthesis
from dashscope.api_entities.dashscope_response import ImageSynthesisResponse
import dashscope # 引入 dashscope 主模块

# --- 加载环境变量 ---
load_dotenv() # 仍然需要
api_key = os.getenv("DASHSCOPE_API_KEY")
if api_key:
    dashscope.api_key = api_key
    print("--- DashScope API Key 已被显式设置！---")
else:
    print("--- 警告：未在环境变量中找到 DASHSCOPE_API_KEY！---")


# --- 初始化组件  ---
PERSIST_DIRECTORY = "../chroma_db"

# 新的代码:
# 使用OllamaEmbeddings
embeddings = OllamaEmbeddings(model="nomic-embed-text")
# 使用ChatOllama，并指定使用我们拉取的llama3模型
llm = ChatOllama(model="llama3", temperature=0)

# 向量数据库的初始化现在使用新的 embeddings 对象
vectordb = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)


# --- RAG 问答链 ---

# 1. 定义 Prompt 模板
prompt_template = """
你是一个世界级的软件工程师和代码分析专家。
你的任务是基于我提供的【代码上下文】，精准地回答我的【问题】。

请遵循以下规则：
1. 你的回答必须完全基于提供的上下文，不要编造信息。
2. 如果上下文中没有足够信息来回答问题，请明确指出“根据提供的上下文，我无法回答这个问题”。
3. 你的回答应该清晰、简洁，并使用Markdown格式化代码（如果需要）。
4. 你的回答一定要是中文的

【代码上下文】:
{context}

【问题】:
{question}

你的专业回答:
"""
PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

# 2. 创建 LLMChain
rag_chain = LLMChain(llm=llm, prompt=PROMPT)


def get_rag_answer(query: str) -> str:
    """
    接收用户问题，执行RAG流程，并返回答案。
    """
    print(f"接收到查询: {query}")

    # 1. 从向量数据库检索相关文档
    # similarity_search 会返回一个文档列表
    retrieved_docs = vectordb.similarity_search(query, k=5)  # k=5表示返回5个最相关的结果

    # 2. 将检索到的文档内容拼接成一个字符串
    context = "\n\n---\n\n".join([doc.page_content for doc in retrieved_docs])

    print("--- 检索到的上下文已准备好 ---")

    # 3. 将上下文和问题传入链中，获得答案
    answer = rag_chain.run({"context": context, "question": query})

    return answer


# --- 图像生成功能 ---

# --- 图像生成功能 (这是主要修改的部分) ---

print(dashscope.api_key)

def generate_logo(prompt: str) -> str:
    """
    使用阿里云通义万相生成图像并返回图像 URL。
    """
    # dashscope.api_key = "sk-4861835d021349748fa037974f3bc5f2"
    print(f"接收到图像生成 Prompt: {prompt}")
    try:
        # 调用通义万相的文生图服务
        response: ImageSynthesisResponse = ImageSynthesis.call(
            model=ImageSynthesis.Models.wanx_v1, # 使用通义万相V1模型
            prompt=prompt,
            n=1,  # 生成1张图片
            size='1024*1024' # 图片尺寸
        )

        # 检查API调用是否成功
        if response.status_code == HTTPStatus.OK:
            # 获取并返回第一张图片的URL
            image_url = response.output.results[0].url
            print("--- 图像生成成功 ---")
            return image_url
        else:
            # 如果失败，打印并返回错误信息
            error_message = f"图像生成失败: Code={response.code}, Message={response.message}"
            print(error_message)
            return error_message

    except Exception as e:
        print(f"调用SDK时发生异常: {e}")
        return "图像生成服务调用异常，请检查网络或依赖配置。"